from constants import *
from dict import *
from pinyins import *

_single_dict = DICT_PINYIN_SINGLE_LITE
_phrase_dict = DICT_PINYIN_PHRASE_BASE
_version = "'2309.1'"


class Unipy(object):
    """
    统一的 Rime 词典文件生成器
    """
    TABLE_DICT = 'unipy_table'
    SCRIPT_DICT = 'unipy_script'

    def __init__(self, aux_file, pinyin_conf=None) -> None:
        super().__init__()
        self.pinyinSingleDict = RimeDict(_single_dict)
        self.pinyinPhraseDict = RimeDict(_phrase_dict)
        self.auxDict = RimeDict(aux_file)  # 辅助码词典
        self.encoder = Encoder() if pinyin_conf is None else Encoder(pinyin_conf)

    def writeTableDict(self, dict_name=TABLE_DICT, import_tables=None):
        """
        生成 Rime 中 table_translator 适用的词库
        :param dict_name: 词典名
        :param import_tables: 导入的其它词典
        :return: self
        """
        """
        生成 Rime 中 table_translator 适用的词库
        :return: self
        """
        print('生成单字词库')
        # 加载基本词库
        self.pinyinSingleDict.load()
        self.auxDict.load()

        # 添加辅助码
        items = []
        for item in self.pinyinSingleDict.items:
            head = self.encoder.encode(item.code)
            rear = self.auxDict.getItem(item.text).code
            items.append(RimeDictItem(item.text, head + rear, item.weight))

        # 排序
        items.sort(key=lambda x: (-x.weight, len(x.code), x.code))

        # 写词典文件
        out = create_utf8_file('../out/{}.dict.yaml'.format(dict_name))
        write_dict_head(out, name=dict_name, version=_version, import_tables=import_tables)
        out.write('...\n')
        out.write('# 单字\n')
        codes = set()
        for item in items:
            short = False  # 是否生成过简码
            for i in range(2, len(item.code)):
                s = item.code[0:i]  # 简码
                if s not in codes:
                    short = True
                    out.write('{}\t{}\t{}\n'.format(item.text, s, item.weight))
                    codes.add(s)
                    break
            if short:  # 简码存在
                out.write('{}\t{}\t{}\n'.format(item.text, item.code, item.weight // 2))  # 有简码时，全码 weight 调低
            else:
                codes.add(item.code)
                out.write('{}\n'.format(item))  # 出简也出全
            out.flush()
        out.close()
        return self

    def writeScriptDict(self):
        """
        生成 Rime 中 script_translator 适用的词库
        :return:
        """
        print('生成拼音词库')
        # 加载基本词库
        self.pinyinSingleDict.load()
        self.pinyinPhraseDict.load()

        # 对拼音编码
        items = []
        for item in self.pinyinSingleDict.items:
            code = self.encoder.encode(item.code)
            items.append(RimeDictItem(item.text, code, item.weight))
        for item in self.pinyinPhraseDict.items:
            if len(item.text) > 7:  # 太长的词组就不要了
                continue
            code = self.encoder.encode(item.code.split(' '))
            if code:
                items.append(RimeDictItem(item.text, code, item.weight))
            else:
                print('不能转换的词条: {}'.format(item))

        # 排序
        items.sort(key=lambda x: (len(x.text), x.code, -x.weight))

        # 写词典文件
        out = create_utf8_file('../out/{}.dict.yaml'.format(Unipy.SCRIPT_DICT))
        write_dict_head(out, name=Unipy.SCRIPT_DICT, version=_version)
        out.write('...\n')
        for item in items:
            out.write('{}\n'.format(item))
            out.flush()
        out.close()
        return self


def _get_full_pinyin(text, pinyin, single_dict):
    if text in single_dict:
        for full in single_dict.get(text):
            if full.startswith(pinyin) and len(pinyin) + 1 == len(full):
                return full
    return pinyin
